#pragma once

#include "NvInferPlugin.h"
#include "cuda_runtime_api.h"

#include <iostream>
#include <memory>
#include <opencv2/core.hpp>
#include <rclcpp/rclcpp.hpp>
#include "../trt_infer.h"

namespace perception::camera {

class LaneSegLogger : public nvinfer1::ILogger {
  public:
    explicit LaneSegLogger(Severity severity = Severity::kWARNING)
        : reportable_severity(severity) {}

    void log(Severity severity, const char* msg) noexcept {
        if (severity > reportable_severity) {
            return;
        }
        switch (severity) {
            case Severity::kINTERNAL_ERROR:
                std::cerr << "lane_seg INTERNAL_ERROR: ";
                break;
            case Severity::kERROR:
                std::cerr << "lane_seg ERROR: ";
                break;
            case Severity::kWARNING:
                std::cerr << "lane_seg WARNING: ";
                break;
            case Severity::kINFO:
                std::cerr << "lane_seg INFO: ";
                break;
            default:
                std::cerr << "lane_seg UNKNOWN: ";
                break;
        }
        std::cerr << msg << std::endl;
    }
    Severity reportable_severity;
};

class LaneSeg {
  public:
    bool init(const std::string& model_dir);
    void deinit();
    void process(const cv::Mat& raw_image, std::vector<std::vector<cv::Point2f>> &contours);
    cv::Mat resize_keep_ratio(const cv::Mat &input, const cv::Size &dstSize, const cv::Scalar &bgcolor);
    void map2mask(cv::Mat real_out,cv::Mat &real_out_);
    static std::shared_ptr<LaneSeg> get_detector(const std::string& model_dir);
    void copy_from_mat(const cv::Mat& image);

  private:
    LaneSeg();
    // ~LaneSeg();
    void LoadEngine(const std::string& engine_file);
    void make_pipe(bool warmup=true);
    void infer();
    void get_mask(cv::Mat lane_mask);



    // void PostProcess(std::vector<std::vector<cv::Point2f>> &lanes, float conf_thresh=0.4f, float nms_thresh=50.f, int nms_topk=4);
    // float LaneIoU(const Detection& a, const Detection& b);
    LaneSegLogger g_logger_;
    rclcpp::Logger logger__;

    // cudaStream_t stream_;
    // nvinfer1::ICudaEngine* engine_;
    // nvinfer1::IExecutionContext* context_;

    nvinfer1::ICudaEngine* engine_ = nullptr;
    nvinfer1::IRuntime* runtime_ = nullptr;
    nvinfer1::IExecutionContext* context_ = nullptr;
    cudaStream_t stream_ = nullptr;

    int num_bindings;
    int num_inputs = 0;
    int num_outputs = 0;
    
    void* buffers_[2];
    int buffer_size_[2];
    std::vector<float> image_data_;
    // std::vector<Detection> detections_;
    static std::shared_ptr<LaneSeg> detector_;
    int num_class_ = 5; // 背景，虚线，实线，斑马区，导流区
    cv::Mat lane_mask_;
    const int infer_h_ = 1080; 
    const int infer_w_ = 1920;
    std::vector<Binding> input_bindings;
    std::vector<Binding> output_bindings;
    std::vector<void*> host_ptrs;
    std::vector<void*> device_ptrs;
    float* data_;
    float *net_out_;
    cv::Mat image_orig_;
    cv::Mat score_map_;
    int infer_count_ =  0;
    std::string input_name_{"input_0"};
    std::string output_name_{"output_0"};
    int pad_top_;
    int pad_left_;
    // void *buffers_[2];



    // int orig_img_h_ = 1080;
    // int orig_img_w_ = 1920;
    // const int INPUT_H_ = 360;
    // const int INPUT_W_ = 640;
    // const int N_OFFSETS_ = 72;
    // const int N_STRIPS_ = (72 - 1);
    // const int MAX_COL_BLOCKS_ = 1000;
};
}